#!/usr/bin/env python

# Imports:
# -------
from config import *
import roslib; roslib.load_manifest(PKG)
import rosbag, log, wx, cv, cv_bridge
from StringIO import StringIO
from annotations.msg import Annotation, AnnotationArray
from gui.main_frame import *
from image_geometry.cameramodels import PinholeCameraModel 
from bag_helpers import *

# ----------------
class BagLabeler():
# ----------------
	""" The guts of the bag labeling tool """

	# ----------------
	def __init__(self):
	# ----------------
		log.info('Initializing the bag labeler')

		# Initialize everything
		self.bagPath = None
		self.bag = None
		self.msgTimes = []
		self.annotations = []
		self.annotationTopic = None
		self.imageTopic = None
		self.rectifyFlag = True
		self.shuffleFlag = False
		self.annotationOutPath = None
		self.unsavedWork = False

		# Create our CV bridge
		self.bridge = cv_bridge.CvBridge()

	# --------------------------------
	def unloadBag(self, callback=None):
	# --------------------------------
		if self.bag != None:
			log.info('Started unloading bag: %s' % self.bagPath)
			self.bag.close()
			log.info('Finished unloading bag: %s' % self.bagPath)

			if callback != None:
				callback()

	# -----------------------
	def analyzeBag(self, bag):
	# -----------------------
		# Extract the list of topics for the relevant types
		log.info('Extracting topic information.')
		topicsByDatatype = get_topics_by_datatype(bag)
		topicsByDatatype.setdefault(ANNOTATION_TOPIC_TYPE, [])

		imageTopics = []
		if IMAGE_TOPIC_TYPE in topicsByDatatype:
			imageTopics = topicsByDatatype[IMAGE_TOPIC_TYPE]

		annotationTopics = []
		if ANNOTATION_TOPIC_TYPE in topicsByDatatype:
			annotationTopics = topicsByDatatype[ANNOTATION_TOPIC_TYPE]

		cameraInfoTopics = []
		if CAMERA_INFO_TOPIC_TYPE in topicsByDatatype:
			cameraInfoTopics = topicsByDatatype[CAMERA_INFO_TOPIC_TYPE]

		return (imageTopics, annotationTopics, cameraInfoTopics)

	# ---------------------------------------
	def loadBag(self, bagPath, callback=None):
	# ---------------------------------------
		# Unload the current bag (if we have one)
		self.unloadBag()

		# Store the bag path
		self.bagPath = bagPath
		log.info('Started loading ROS bag: %s' % bagPath)

		# Load the bag
		self.bag = rosbag.Bag(bagPath, 'r')

		# Analyze the bag
		imageTopics, annotationTopics, cameraInfoTopics = self.analyzeBag(self.bag)
		
		# Execute the callback
		if callback != None:
			if not callback(imageTopics, annotationTopics, cameraInfoTopics):
				# The callback aborted the loading process
				log.info('Bag loading aborted: %s' % bagPath)
				self.bag.close()
				return # Done!
		
		log.info('Finished loading ROS bag: %s' % bagPath)
		log.info('  + Image Topic: %s' % self.imageTopic)
		log.info('  + Annotation Topic: %s' % self.annotationTopic)
		log.info('  + Perform Rectification: %s' % str(self.rectifyFlag))
		log.info('  + Shuffle images: %s' % str(self.shuffleFlag))

	# ----------------------------------------------------------------------------------------------
	def setImageTopic(self, imageTopic, rectifyFlag, shuffleFlag, cameraInfoTopic, startTime = None):
	# ----------------------------------------------------------------------------------------------
		log.info('Matched camera info topic: ' + cameraInfoTopic)

		self.rectifyFlag = rectifyFlag
		self.shuffleFlag = shuffleFlag
		self.imageTopic = imageTopic
		self.cameraInfoTopic = cameraInfoTopic
		self.msgTimes = get_message_times(self.bag, imageTopic, shuffleFlag)

		# Set the current time & index
		startIndex = 0
		if startTime != None:
			for i, time in enumerate(msgTimes):
				if startTime < time:
					startIndex = i
					break
		self.index = startIndex

		# If we're rectifying then we need to load a camera model
		if rectifyFlag:
			log.info('Loading camera model for rectification')
			self.model = PinholeCameraModel()
			topic, msg, t = self.getCameraInfo()
			self.model.fromCameraInfo(msg)	

	# ------------------------------------------------------------------------
	def setAnnotationTopic(self, annotationTopic, bag=None, shuffleFlag=False):
	# ------------------------------------------------------------------------
		log.info('Setting annotation topic: %s' % str(annotationTopic))
		bag = self.bag if bag == None else bag
		self.annotationTopic = annotationTopic

		# Create the empty annotation list
		self.annotations = [None for t in self.msgTimes]

		# Read in the annotations
		for out in bag.read_messages(topics = [annotationTopic]):
			a = out[1]

			# This is UGLY and slow... sorry?
			try:			
				i = self.msgTimes.index(a.header.stamp)
				self.annotations[i] = a
			except ValueError:
				log.error('Unable to load anntotation topic')

	# -------------------------------------------
	def newAnnotationTopic(self, annotationTopic):
	# -------------------------------------------
		log.info('Creating a new annotation topic') 
		self.annotationTopic = annotationTopic
		
		# We're going to create an empty annotation array
		self.annotations = [None for t in self.msgTimes]

		log.info('Done creating new annotation topic')

	# --------------------------------------------------
	def loadExternalAnnotations(self, path, shuffleFlag):
	# --------------------------------------------------
		bag = rosbag.Bag(path, 'r')

		try:
			imageTopics, annotationTopics, cameraInfoTopics = self.analyzeBag(bag)
			if len(annotationTopics) > 0:
				self.setAnnotationTopic(annotationTopics[0], bag, shuffleFlag)
			else:
				return False

		finally:
			bag.close()

		return True

	# -----------------------------------
	def convertToCvImage(self, rosImgMsg):
	# -----------------------------------
		try:
			# Convert the image message to a format OpenCV plays nice with
			width =  rosImgMsg.width
			height = rosImgMsg.height
			cvImg = self.bridge.imgmsg_to_cv(rosImgMsg, 'bgr8') # TODO: Conversion should be better/safer...
			return (width, height, cvImg)

		except Exception as e: 
			log.error('Unable to convert ros image message successfully to OpenCV image: %s' % str(e))
			return None


	# ---------------------------------------------
	def convertToBitmap(self, width, height, cvImg):
	# ---------------------------------------------
		try:
			# Convert the image to something WX can display
			wxImage = wx.EmptyImage(width, height)
			wxImage.SetData(cvImg.tostring())
			return wxImage.ConvertToBitmap()

		except AttributeError as e: 
			return None

	# ---------------------
	def getCameraInfo(self):
	# ---------------------
		# Get the image
		gen = self.bag.read_messages(topics=[self.cameraInfoTopic])
		topic, msg, t = gen.next()
		
		return (topic, msg, t)


	# -----------------------
	def getImage(self, index):
	# -----------------------
		# Get the image
		startTime = self.msgTimes[index]
		gen = self.bag.read_messages(topics=[self.imageTopic], start_time=startTime)
		topic, msg, t = gen.next()
		
		return (topic, msg, t)

	# -----------------------------------------
	def getImageAndAnnotationArray(self, index):
	# -----------------------------------------
		# Get the image
		topic, msg, t = self.getImage(index)

		# Convert the image to OpenCV's format
		width, height, cvImg = self.convertToCvImage(msg)

		# Perform rectification if required
		if self.rectifyFlag:
			log.info('Rectifying image #%i' % index) 

			# Create a header for the rectified image
			rectifiedImg = cv.CreateMat(cvImg.rows, cvImg.cols, cv.CV_8UC3)

			# Perform the rectification
			self.model.rectifyImage(cvImg, rectifiedImg)
		
			# Replace the origional image
			cvImg = rectifiedImg

		# Convert the image to a bitmap
		bmp = self.convertToBitmap(width, height, cvImg)		
	
		# Get the annotation array (if we have one)
		annotationArray = None
		if len(self.annotations) > 0:
			# If we don't have an annotation array at the index, make one
			if self.annotations[index] == None:
				aa = AnnotationArray()
				aa.header = msg.header
				self.annotations[index] = aa
			
			annotationArray = self.annotations[index]

		return topic, bmp, t, annotationArray

	# -----------------------
	def setIndex(self, index):
	# -----------------------
		self.index = index
		return self.getImageAndAnnotationArray(index)

	# ----------------
	def getIndex(self):
	# ----------------
		return (self.index, len(self.msgTimes))

	# ------------
	def curr(self):
	# ------------
		return self.getImageAndAnnotationArray(self.index)

	# ------------
	def next(self):
	# ------------		
		# If the annotation array is empty, remove it before moving on
		if self.annotations[self.index] != None and len(self.annotations[self.index].annotations) == 0:
			self.annotations[self.index] = None

		# Increment the index
		self.index += 1
		if self.index >= len(self.msgTimes):
			self.index = 0

		return self.getImageAndAnnotationArray(self.index)

	# ------------
	def prev(self):
	# ------------
		# If the annotation array is empty, remove it before moving on
		if self.annotations[self.index] != None and len(self.annotations[self.index].annotations) == 0:
			self.annotations[self.index] = None

		# Decrement the index
		self.index -= 1
		if self.index < 0:
			self.index = len(self.msgTimes) - 1

		return self.getImageAndAnnotationArray(self.index)

	# ------------------
	def isBagLoaded(self):
	# ------------------
		return self.bag != None

	# ----------------
	def shutdown(self):
	# ----------------
		log.info('Shutting down the bag labeler')
		self.unloadBag()

	# -----------------------------
	def saveAnnotations(self, path):
	# -----------------------------
		log.info('Saving annotations to \'%s\'' % path)

		# Open the out bag		
		outBag = rosbag.Bag(path, 'w')	

		try:
			# Loop through the annotations and save		
			annotationTopic = self.annotationTopic
			for annotationArray in self.annotations:

				# Only write non-null annotations
				if annotationArray != None:

					"""
					# Strip out invalid annotations (but leave the array to indicate nothing was found)
					for i, annotation in enumerate(annotationArray.annotations):
						if annotation.x < 0 or annotation.y < 0 or annotation.width < 0 or annotation.height < 0:
							del annotationArray.annotations[i]
					"""		

					outBag.write(annotationTopic, annotationArray, annotationArray.header.stamp)
				
		finally:
			outBag.close()

		self.unsavedWork = False
		log.info('Done saving annotations')

# ---------------------
if __name__=='__main__':
# ---------------------
	# Create the labeler
	bagLabeler = BagLabeler()

	# Create the GUI
	app = wx.App()
	frame = MainFrame(bagLabeler, WINDOW_TITLE, WINDOW_SIZE)
	frame.Show()
	app.MainLoop()
